
#pragma once

#include <future>
#include <thread>
#include <chrono>

#include "nonstd/scope.hpp"
#include "SyncBox.hpp"
#include "latch.hpp"

#ifdef _MSC_VER
#include <windows.h>
#endif   // _MSC_VER

#include "BlockingCollection.h"
using namespace code_machina;

#include "task_thread_pool.hpp"

template<class _Tx, class _Ty>
class Box
{
public:
    void sync_wait(long                             ms,
                   const std::function<void()>&     on_send    = nullptr,
                   const std::function<void(_Ty&)>& on_receive = nullptr,
                   const std::function<void()>&     on_timeout = nullptr)
    {
        outbox_pr_ = std::make_unique<std::promise<_Ty>>();
        auto f = outbox_pr_->get_future();

        auto _ = std::async(std::launch::async, [&]() {
            if (on_send != nullptr)
                on_send();
        });

        alive_.store(true);
        auto status = f.wait_for(std::chrono::milliseconds(ms));

        if (status == std::future_status::ready)
        {
            if (on_receive != nullptr)
            {
                _Ty vec = std::move(f.get());
                on_receive(vec);
            }
        }
        else
        {
            alive_.store(false);
            if (on_timeout != nullptr)
                on_timeout();
        }

        box_wait_fininsh();
    }

    void async_notify(_Ty& x)
    {
        if (alive_.load())
        {
            alive_.store(false);

            outbox_pr_->set_value(std::move(x));
        }
        box_notify();
    }

    void box_used()
    {
        std::unique_lock<std::mutex> lock(mtx);
        used = true;
    }

private:
    void box_wait_fininsh()
    {
        std::unique_lock<std::mutex> lock(mtx);
        if (used)
        {
            cv.wait(lock);
        }
    }

    void box_notify()
    {
        std::unique_lock<std::mutex> lock(mtx);
        used = false;
        cv.notify_one();
    }

public:
    _Tx inbox_;

private:
    std::mutex              mtx;
    std::condition_variable cv;
    bool                    used{false};

    std::unique_ptr<std::promise<_Ty>> outbox_pr_;
    std::atomic_bool                   alive_{false};
};

template<class _Tx, class _Ty, size_t N = 1>
class Sync
{
public:
    Sync(std::function<_Ty(_Tx&)>& do_work)
    {
        std::thread([&]() {
            while (true)
            {
                uint64_t idx;
                ids_.take(idx);
                if (exit_)
                {
                    break;
                }

                pool.submit(
                    [=](uint64_t idx) {
                        auto box = box_get(idx, [](Box<_Tx, _Ty>* box) {
                            box->box_used();
                        });
                        if (box)
                        {
                            _Ty tout = do_work(box->inbox_);

                            box->async_notify(tout);
                        }
                    },
                    idx);
            }
        }).detach();
    }

    ~Sync()
    {
        exit_.store(true);
        ids_.emplace(0);
    }

    void transmit(long                                   ms,
                  const std::function<_Tx()>&            on_send    = nullptr,
                  const std::function<void(_Tx&, _Ty&)>& on_receive = nullptr,
                  const std::function<void(_Tx&)>&       on_timeout = nullptr)
    {
        uint64_t idx = index_++;

        //std::this_thread::sleep_for(std::chrono::milliseconds(1000));
        auto box            = new Box<_Tx, _Ty>();
        auto sync_box_guard = nonstd::make_scope_exit([&]() {
            delete box;
        });

        auto on_send_ = [&]() {
            box->inbox_ = std::move(on_send());

            std::lock_guard<std::mutex> lck(mtx_);

            map_.insert(std::make_pair(idx, box));

            ids_.emplace(idx);
        };
        auto on_timeout_ = [&]() {
            box_erase(idx);

            on_timeout(box->inbox_);
        };
        auto on_receive_ = [&](_Ty& TOut) {
            box_erase(idx);

            on_receive(box->inbox_, TOut);
        };

        box->sync_wait(ms, on_send_, on_receive_, on_timeout_);
    }

private:
    Box<_Tx, _Ty>* box_get(uint64_t idx, const std::function<void(Box<_Tx, _Ty>*)>& fn = nullptr)
    {
        std::lock_guard<std::mutex> lck(mtx_);

        Box<_Tx, _Ty>* ptr = NULL;

        try
        {
            ptr = map_.at(idx);

            fn(ptr);
        }
        catch (const std::exception&)
        {
        }
        return ptr;
    }

    void box_erase(uint64_t idx)
    {
        std::lock_guard<std::mutex> lck(mtx_);
        map_.erase(idx);
    }

private:
    task_thread_pool::task_thread_pool pool{N};

    std::mutex                                   mtx_;
    std::atomic<uint64_t>                        index_{0};
    std::unordered_map<uint64_t, Box<_Tx, _Ty>*> map_;
    BlockingQueue<uint64_t>                      ids_;
    std::atomic_bool                             exit_{false};
};
